import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.path.append('/root/ptv3/PointTransformerV3/Pointcept/pointcept/utils')
from visualization import *
import open3d as o3d
import numpy as np

S3DIS_COLOR_MAP_13 = {
    0: (233.0, 228.0, 107.0),
    1: (113.0, 186.0, 235.0),
    2: (214.0, 138.0, 96.0),
    3: (240.0, 148.0, 136.0),
    4: (255.0, 61.0, 61.0),
    5: (77.0, 173.0, 83.0),
    6: (108.0, 134.0, 74.0),
    7: (48.0, 58.0, 121.0),
    8: (107.0, 56.0, 113.0),
    9: (96.0, 196.0, 117.0),
    10: (95.0, 95.0, 91.0),
    11: (239.0, 239.0, 239.0),
    12: (96.0, 131.0, 136.0),
}

def get_npy_files(folder_path):
    """
    获取指定文件夹下所有.npy文件的路径列表
    
    参数:
        folder_path (str): 目标文件夹路径
    
    返回:
        list: 包含所有.npy文件路径的列表
    """
    npy_files = []
    
    # 检查文件夹是否存在
    if not os.path.exists(folder_path):
        print(f"错误: 文件夹 {folder_path} 不存在")
        return npy_files
    
    # 遍历文件夹及其子文件夹
    for root, dirs, files in os.walk(folder_path):
        for file in files:
            # 检查文件扩展名是否为.npy
            if file.endswith('.npy'):
                npy_files.append(os.path.join(root, file))
    
    return npy_files

def save_point_cloud(coord, color, filename):
    """
    Save point cloud data to a PLY file.
    
    Parameters:
    - coord: numpy array of shape (N, 3) representing the coordinates of the points.
    - color: numpy array of shape (N, 3) representing the colors of the points.
    - filename: string, the name of the file to save the point cloud.
    """
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(coord)
    pcd.colors = o3d.utility.Vector3dVector(color)
    o3d.io.write_point_cloud(filename, pcd)

def base_test():
    coord = np.load("/root/ptv3/PointTransformerV3/Pointcept/data/s3dis/Area_5/conferenceRoom_1/coord.npy")
    # color = np.load("/root/ptv3/PointTransformerV3/Pointcept/data/s3dis/Area_5/conferenceRoom_1/color.npy")
    # color = color.astype(np.float64)/255  # 确保颜色是 float64 类型
    segment = np.load("/root/ptv3/PointTransformerV3/Pointcept/data/s3dis/Area_5/conferenceRoom_1/segment.npy")
    pred = np.load("/root/ptv3/PointTransformerV3/Pointcept/exp/s3dis/semseg-pt-v3m1-0-base/result/Area_5-conferenceRoom_1_pred.npy")
    # instance = np.load("/root/ptv3/PointTransformerV3/Pointcept/data/s3dis/Area_5/conferenceRoom_1/instance.npy")
    color_seg = np.array([S3DIS_COLOR_MAP_13[i.item()] for i in segment], dtype=np.float64) / 255.0  # 确保颜色是 float64 类型
    color_pred = np.array([S3DIS_COLOR_MAP_13[i.item()] for i in pred], dtype=np.float64) / 255.0  # 确保颜色是 float64 类型
    # print(color.shape, color.dtype)  # 应输出类似 (N,3) float64
    # print(set(segment.flatten()))  # 打印前10个点的颜色

    save_point_cloud(coord, color_pred, "pred.ply")
    save_point_cloud(coord, color_seg, "seg.ply")


def visual_process():
    pred_folder = "/root/ptv3/PointTransformerV3/Pointcept/exp/s3dis/semseg-pt-v3m1-0-base/result"
    pred_files = get_npy_files(pred_folder)
    for pred_file in pred_files:
        import re
        # 正则表达式匹配模式：
        # - 匹配 "Area_X-" 后的内容，直到 "_pred" 前
        pattern = r"Area_\d+-(.*)_pred\.npy"
        match = re.search(pattern, pred_file)

        if match:
            room = match.group(1)  # 输出 "conferenceRoom_1"
            print(room)
        else:
            print("未匹配到目标名称")

        coord_file = f"/root/ptv3/PointTransformerV3/Pointcept/data/s3dis/Area_5/{room}/coord.npy"
        # segment_file = f"/root/ptv3/PointTransformerV3/Pointcept/data/s3dis/{area}/{room}/segment.npy"

        if not os.path.exists(coord_file):
            print(f"Missing files for {room}, skipping...")
            continue

        coord = np.load(coord_file)
        # segment = np.load(segment_file)
        pred = np.load(pred_file)

        # color_seg = np.array([S3DIS_COLOR_MAP_20[i.item()] for i in segment], dtype=np.float64) / 255.0
        color_pred = np.array([S3DIS_COLOR_MAP_13[i.item()] for i in pred], dtype=np.float64) / 255.0

        save_point_cloud(coord, color_pred, f"s3dis_Area5_ply/{room}_pred.ply")
        # save_point_cloud(coord, color_seg, f"{area}_{room}_seg.ply")


if __name__ == "__main__":
    # save_point_cloud(coord, color, "test.ply")
    visual_process()
    print("Point cloud saved successfully.")